[WS1][kernels] Add fixed-order joint-attention softmax for Qwen-Image - #453

Latest Status [8 Oct 2026]

• PyTorch reference forward/backward
• CUDA forward/backward
• Triton forward/backward
• Registry, bindings, and backend fallback
• Correctness and bitwise tests on RTX 5060
• Documentation and benchmark script
• Compare fixed tile sizes (64/128/256/512)
• Add benchmark results for the selected tile size
• Add joint_attn_softmax to the unified gtest
• Check prompt padding and record the actual backend

Summary

Following the Qwen-Image WS1 attention breakdown in #386, this PR implements
joint_attn_softmax between the QK GEMM and AV GEMM:

joint_attn_qk_gemm -> joint_attn_softmax -> joint_attn_av_gemm

The operator takes the attention scores produced by QK GEMM, applies softmax
along the key dimension, and passes the probabilities to AV GEMM.

The main point here is not just to implement softmax, but also to fix its
floating-point calculation and reduction order. The same row must not produce
different results because of its batch, launch geometry, or selected backend.

This PR includes:

• Fixed-order online softmax

  • One CUDA block or Triton program owns one complete row.
  • The key dimension is split into fixed-size tiles and processed from left to
    right.
  • Each tile uses a fixed reduction tree for max and sum, and the online
    softmax states are merged in a fixed order.
  • Intermediate calculations use FP32, with one dtype conversion at the final
    output or gradient write.
  • Split-K, Stream-K, atomics, TF32, fast math, and autotuning are not used.

• PyTorch, CUDA, and Triton backends

  • PyTorch defines the calculation and reduction order as the CPU reference.
  • CUDA is the NVIDIA GPU implementation and the bit-level reference backend.
  • Triton uses the same tile order, reduction tree, and exponential
    approximation as CUDA.
  • All three backends support forward and backward.

• Fixed-order backward

  • Backward uses the standard softmax VJP:

    delta = fixed_sum(P * dP)
    dScores = P * (dP - delta)
    

  • fixed_sum(P * dP) uses the same tile layout and reduction order as
    forward.

  • Gradients are calculated in FP32 and converted to the input dtype at the
    final write.

• Registry and fallback

  • On NVIDIA GPUs, the selection order is CUDA -> Triton -> PyTorch.
  • If the CUDA extension is unavailable or was built with an incompatible
    fast-math configuration, dispatch falls through explicitly.
  • The trace records the selected backend, reduction order, accumulator
    precision, and kernel fingerprint.

I compared fixed tile sizes 64/128/256/512 on an RTX 5060. The best size
varies by backend, but 256 gives a reasonable balance across CUDA and Triton,
so this PR keeps TILE_K=256 fixed.

Files

Test

python -m pytest -p no:cacheprovider \
  tests/test_joint_attn_softmax.py \
  tests/test_joint_attn_softmax_cuda.py \
  tests/test_joint_attn_softmax_triton.py \
  tests/test_joint_attn_softmax_registry.py \
  tests/test_joint_attn_softmax_full_shapes.py \
  tests/test_build_platform_collectives.py \
  tests/test_operator_inputs.py -q

I also ran the unified gtest for all three backends:

python scripts/check_operator.py --op joint_attn_softmax --candidate pytorch \
  --device cpu --dtype fp32 --batch 2 --seq 257 --check-grad
python scripts/check_operator.py --op joint_attn_softmax --candidate cuda \
  --device cuda --dtype bf16 --batch 2 --seq 257 --check-grad
python scripts/check_operator.py --op joint_attn_softmax --candidate triton \
  --device cuda --dtype bf16 --batch 2 --seq 257 --check-grad

Test results

Environment:

Result: 140 passed, 0 skipped, 0 failed.

Benchmark

The benchmark checks bytes against CUDA before timing:

python benchmarks/benchmark_joint_attn_softmax.py --dtype bf16
python benchmarks/benchmark_joint_attn_softmax.py --dtype bf16 --backward
python benchmarks/benchmark_joint_attn_softmax.py --dtype fp32
python benchmarks/benchmark_joint_attn_softmax.py --dtype fp32 --backward

--backward measures forward plus torch.autograd.grad, not the isolated
backward kernel.

I first compared fixed TILE_K=64/128/256/512 builds. The sweep used 16,384
rows, the three reference key lengths, BF16 and FP32, 10 warmups and 50 measured
iterations. Each value below is the average of the three shape medians across
three runs. Times are forward + backward kernel latency in milliseconds.

128 was slightly better for the main CUDA BF16 path, while 256 was better
for CUDA FP32. Triton preferred a larger tile, but 512 regressed CUDA. Since
the three backends share one arithmetic contract, this PR keeps TILE_K=256 as
the common choice.

All 36 forward/backward checks per GPU backend matched the PyTorch reference for
the corresponding tile. Different arithmetic tile sizes were not byte-equal to
each other, so the selected tile remains fixed and runtime autotuning is not
used.

Selected TILE_K=256 BF16 results (rows=24, median of five runs, five warmups
and 20 measured iterations):

PyTorch, CUDA and Triton were byte-equal in the selected-tile benchmark.